package com.mantou.mantouaiagent.rag.queryAugmenter;


import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;

import java.util.HashMap;
import java.util.Map;

/**
 * 通用用户上下文处理增强工厂
 */
public class GeneralContextQueryAugmenterFactory {

    /**
     * 不允许生成别的上下文
     * @param type
     * @return
     */
    public static ContextualQueryAugmenter createInstanceFalse(String type)
    {
    // 定义带有变量的模板
    String template = """
                你应该输出下面的内容：
                抱歉，我只能回答{type}的问题，请重新提问哦 避免询问无关问题
                """;

    PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("type",type));

        return ContextualQueryAugmenter.builder()
                //不允许生成别的上下文 false
                .allowEmptyContext(false)
                .emptyContextPromptTemplate(promptTemplate)
                .build();
    }
    /**
     * 不允许生成别的上下文
     * @param
     * @return
     */
    public static ContextualQueryAugmenter createInstanceTrue()
    {
        return ContextualQueryAugmenter.builder()
                //不允许生成别的上下文
                .allowEmptyContext(true)
//                .emptyContextPromptTemplate(promptTemplate)
                .build();
    }
}
